package binary_tree

func IsBalanced(root *Node) bool {
	if root == nil {return true}
	ans := []bool{true}
	balanceTraversal(root, ans)
	return ans[0]
}

func balanceTraversal(node *Node, ans []bool) int {
	if !ans[0] || node == nil {
		return 0
	}
	leftHeigt  := balanceTraversal(node.Left, ans)
	rightHeigt := balanceTraversal(node.Right, ans)
	if leftHeigt - rightHeigt > 1 || leftHeigt - rightHeigt < -1 {
		ans[0] = false
	}
	return max(leftHeigt, rightHeigt) + 1
}

func IsBalanced2(root *Node) bool {
	if root == nil {return true}
	return balanceProcess(root).isBalance
}

type BalanceInfo struct {
	isBalance bool
	height    int
}

func balanceProcess(node *Node) BalanceInfo {
	if node == nil {
		return BalanceInfo{isBalance: true, height: 0}
	}
	leftInfo := balanceProcess(node.Left)
	rightInfo := balanceProcess(node.Right)

	var height int
	height = max(leftInfo.height, rightInfo.height) + 1
	isBalance := false
	if abs(leftInfo.height - rightInfo.height) <= 1 && leftInfo.isBalance && rightInfo.isBalance{
		isBalance = true
	}
	return BalanceInfo{isBalance, height}
}

func abs(a int) int {
	if a < 0 {
		return ^a+1
	}
	return a
}

